package com.mybatis.generator.db;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.mybatis.generator.bean.GenColumn;
import com.mybatis.generator.bean.GenTable;
import com.mybatis.generator.config.CodeGeneratorConfig;
import com.mybatis.generator.utils.GenUtil;

import java.sql.*;
import java.util.List;
import java.util.Map;

public class GenDb {

	private CodeGeneratorConfig generatorConfig;

	private Connection conn = null;
	private Map<String, String> javaTypeMap = Maps.newHashMap();
	private Map<String, String> jdbcTypeMap = Maps.newHashMap();

	public GenDb(CodeGeneratorConfig config) {
		this.generatorConfig = config;
		try {
			Class.forName(config.getDbDriverClassName()).newInstance();
			conn = DriverManager.getConnection(config.getDbUrl(), config.getDbUserName(), config.getDbUserPassword());
			// 数据初始化集合对象
			javaTypeMap.put("bigint", "Long");
			javaTypeMap.put("varchar", "String");
			javaTypeMap.put("datetime", "Date");
			javaTypeMap.put("char", "String");
			javaTypeMap.put("tinyint", "Integer");
			javaTypeMap.put("int", "Integer");
			javaTypeMap.put("double", "Double");
			javaTypeMap.put("smallint", "Integer");
			javaTypeMap.put("bigint", "Long");
			javaTypeMap.put("text", "String");
			javaTypeMap.put("date", "Date");
			javaTypeMap.put("bit", "Boolean");
			javaTypeMap.put("decimal", "BigDecimal");

			jdbcTypeMap.put("int", "INTEGER");
			jdbcTypeMap.put("text", "VARCHAR");
			jdbcTypeMap.put("longtext", "LONGVARCHAR");
			jdbcTypeMap.put("date", "TIMESTAMP");
			jdbcTypeMap.put("datetime", "TIMESTAMP");
			jdbcTypeMap.put("double", "NUMERIC");
			jdbcTypeMap.put("bit", "BIT");
			jdbcTypeMap.put("decimal", "DECIMAL");
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * 获取数据库的所有表
	 */
	public List<GenTable> getAllTableName() throws Exception {
		List<GenTable> list = Lists.newArrayList();
		PreparedStatement ps = conn.prepareStatement("select table_name, table_comment from information_schema.tables" + " where table_schema='" + generatorConfig.getDbName() + "'");
		ResultSet rs = ps.executeQuery();
		while (rs.next()) {
			GenTable genTable = new GenTable();
			genTable.setTableName(rs.getString("table_name"));
			genTable.setTableComment(rs.getString("table_comment"));
			list.add(genTable);
		}
		return list;
	}

	/**
	 * 获取单表的字段名称，注释
	 */
	public List<GenColumn> getTableNameColumn(String tableName) throws Exception {
		List<GenColumn> list = Lists.newArrayList();
		PreparedStatement ps = conn.prepareStatement("select distinct column_name,column_comment,data_type from "
				+ "information_schema.columns where table_name='" + tableName + "' and table_schema='" + generatorConfig.getDbName() + "'");
		ResultSet rs = ps.executeQuery();
		while (rs.next()) {
			GenColumn column = new GenColumn();
			column.setColumnName(GenUtil.replaceUnderlineAndfirstToUpper(rs.getString(1), "_", ""));
			column.setDbName(rs.getString(1));
			column.setColumnComment(rs.getString(2));
			column.setDataType(getDataType(rs.getString(3)));
			column.setDbType(getJdbcType(rs.getString(3)));
			list.add(column);
		}
		return list;
	}

	private String getDataType(String dataType) {
		if (javaTypeMap.containsKey(dataType)) {
			return javaTypeMap.get(dataType);
		}
		return "String";
	}

	private String getJdbcType(String dbType) {
		if (jdbcTypeMap.containsKey(dbType)) {
			return jdbcTypeMap.get(dbType);
		}
		return dbType;
	}

}
